import os
import sys

import torch
from zkl_aiutils_datasets import load_dataset
from zkl_llmpt_iterator import QwenTokenizer, Utf8Tokenizer, make_llmpt_dataset_torch

project_dir_path = os.path.join(os.path.dirname(__file__), "../..")
sys.path.append(project_dir_path)


def load_preset_text_tokenizer(text_tokenizer_name: str):
    match text_tokenizer_name:
        case 'utf8':
            return Utf8Tokenizer()
        case 'qwen':
            return QwenTokenizer()
        case _:
            raise ValueError(f"Unsupported tokenizer name: {text_tokenizer_name}")


def load_preset_text_dataset(text_dataset_name: str, split_name: str | None = None):
    if split_name is not None:
        dataset = load_preset_text_dataset(text_dataset_name)
        return dataset.named_children[split_name]

    datasets_dir_path = os.path.join(project_dir_path, "datasets")
    match text_dataset_name:
        case 'wiki':
            return load_dataset(os.path.join(datasets_dir_path, "text/wiki/v1.1"))
        case 'mixture':
            return load_dataset(os.path.join(datasets_dir_path, "text/mixture/v6"))
        case _:
            raise ValueError(f"Unsupported dataset name: {text_dataset_name}")


def load_preset_llmpt_dataset(*,
    text_dataset_name: str,
    text_tokenizer_name: str,
    split_name: str | None,
    batch_samples_n: int,
    chunk_tokens_n: int,
    device: torch.device | str | None = None,
):
    tokenizer = load_preset_text_tokenizer(text_tokenizer_name)
    dataset = load_preset_text_dataset(text_dataset_name, split_name)

    match (text_dataset_name, text_tokenizer_name, split_name):
        case ('wiki', 'utf8', 'train'):
            total_tokens_n = 19.7e9
        case ('wiki', 'qwen', 'train'):
            total_tokens_n = 3.69e9
        case _:
            total_tokens_n = None

    return make_llmpt_dataset_torch(
        dataset=dataset,
        tokenizer=tokenizer,
        batch_samples_n=batch_samples_n,
        chunk_tokens_n=chunk_tokens_n,
        total_tokens_n=total_tokens_n,
        device=device)
